#include "head.cuh"

void simple_sort_desc(float *out,
	const float *in,
	const unsigned int nDim){
	memcpy(out, in, nDim*sizeof(float));
	quicksort(out, 0, nDim);
}

void SimplexPr(float *z,
	const float *v,
	const float* dgv,
	unsigned int nS,
	float b){
	//z = SimplexPr(v - dgv / b);
	float *X = (float*)malloc(nS*sizeof(float));
	for (int i = 0; i < nS; i++){
		X[i] = v[i] - dgv[i] / b;
	}
	float *t = (float *)malloc(nS*sizeof(float));
	simple_sort_desc(t, X, nS);
	int kk = -1;
	float tep;
	float sum_t_1_j;
	for (int j = 0; j < nS; j++){
		sum_t_1_j = 0.0f;
		for (int k = 0; k < j; k++)
			sum_t_1_j += t[k];
		tep = t[j] - (sum_t_1_j - 1) / (j + 1);
		if (tep <= 0){
			kk = j - 1;
			break;
		}
	}
	if (kk == -1)
		kk = nS;
	float theta = 0.0f;
	for (int i = 0; i <= kk; i++)
		theta += t[i];
	theta -= 1;
	theta /= kk + 1;
	for (int i = 0; i < nS; i++){
		z[i] = ((X[i] - theta) > 0) ? (X[i] - theta) : 0;
	}
	free(X);
	free(t);
}

void LAE_process(const float *x, const unsigned int nDim,
	const float *U, const unsigned int nS,
	float *val_one_line,/* out */
	const unsigned int num_of_LAE_iters,
	cublasHandle_t cublashandle){
	float *z0, *z1;
	z0 = (float*)malloc(nS*sizeof(float));
	z1 = (float*)malloc(nS*sizeof(float));
	for (int i = 0; i < nS; i++){z0[i] = 1.0 / nS; z1[i] = 1.0 / nS;}
	float *delta, *beta;
	delta = (float*)malloc((num_of_LAE_iters + 2)*sizeof(float));
	beta = (float*)malloc((num_of_LAE_iters + 1)*sizeof(float));
	for (int i = 0; i < num_of_LAE_iters + 2; i++) delta[i] = 0.0f;
	for (int i = 0; i < num_of_LAE_iters + 1; i++) beta[i] = 0.0f;
	delta[0] = 0; delta[1] = 1; beta[0] = 1;

	float alpha;
	float *v = (float*)malloc(nS*sizeof(float));
	float *dif = (float*)malloc(nDim*sizeof(float));
	float gv, gz, gvz;
	float *dgv = (float*)malloc(nS*sizeof(float));
	float tmp;
	float *s_tmp = (float*)malloc(nS*sizeof(float));
	float b;
	float *z = (float *)malloc(nS*sizeof(float));
	float sum_abs_z1_z0;

	float *d_v, *d_U, *d_x;
	// Note again: U is nS_X_nDim
	float *d_dif,*d_dgv;
	float *d_z;
	CHECK_CUDA(cudaMalloc(&d_U, nS*nDim*sizeof(float)));
	CHECK_CUDA(cudaMemcpy(d_U, U, nS*nDim*sizeof(float),cudaMemcpyHostToDevice));
	CHECK_CUDA(cudaMalloc(&d_v, nS*sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_x, nDim*sizeof(float)));
	CHECK_CUDA(cudaMemcpy(d_x, x, nDim*sizeof(float), cudaMemcpyHostToDevice));
	CHECK_CUDA(cudaMalloc(&d_dif, nDim*sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_dgv, nS*sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_z, nS*sizeof(float)));

	float *d_sXs_tmp;
	float *d_sX1_tmp;
	CHECK_CUDA(cudaMalloc(&d_sXs_tmp, nS*nS*sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_sX1_tmp, nS*sizeof(float)));
	float alphaf;
	float betaf;
	for (int i = 0; i < num_of_LAE_iters; i++){
		alpha = (delta[i] - 1) / delta[i + 1];
		for (int j = 0; j < nS; j++)
			v[j] = z1[j] + alpha*(z1[j] - z0[j]);
		CHECK_CUDA(cudaMemcpy(d_v,v,nS*sizeof(float),cudaMemcpyHostToDevice));

		alphaf = -1.0f;
		betaf = 1.0f;
		CHECK_CUDA(cudaMemcpy(d_dif, x, nDim*sizeof(float), cudaMemcpyHostToDevice));
		CHECK_CUDA_STATUS(cublasSgemv(cublashandle, CUBLAS_OP_T, nS, nDim, &alphaf, d_U, nS, d_v, 1, &betaf, d_dif, 1));
		
		CHECK_CUDA(cudaMemcpy(dif,d_dif,nDim*sizeof(float), cudaMemcpyDeviceToHost));
		tmp = 0.0f;
		for (int j = 0; j < nDim; j++)
			tmp += dif[j] * dif[j];
		gv = tmp / 2;
		
		alphaf = 1.0f;
		betaf = 0.0f;
		CHECK_CUDA_STATUS(cublasSgemm(cublashandle, CUBLAS_OP_N, CUBLAS_OP_T, nS, nS, nDim, &alphaf, 
			d_U, nS, d_U, nS, &betaf, d_sXs_tmp, nS));
		CHECK_CUDA_STATUS(cublasSgemv(cublashandle, CUBLAS_OP_N, nS, nS, &alphaf, d_sXs_tmp, nS, d_v, 1, &betaf, d_dgv, 1));
		alphaf = -1.0f;
		betaf = 1.0f;
		CHECK_CUDA_STATUS(cublasSgemv(cublashandle, CUBLAS_OP_N, nS, nDim, &alphaf, d_U, nS, d_x, 1, &betaf, d_dgv, 1));
		CHECK_CUDA(cudaMemcpy(dgv, d_dgv, nS*sizeof(float),cudaMemcpyDeviceToHost));

		// seek beta
		for (int k = 0; k < 100; k++){
			if (k == 0){
				b = beta[i];
			}
			else{
				b = (2 << (k - 1))*beta[i];
			}
			SimplexPr(z, v, dgv, nS, b);
			alphaf = -1.0f;
			betaf = 1.0f;
			CHECK_CUDA(cudaMemcpy(d_dif, x, nDim*sizeof(float), cudaMemcpyHostToDevice));
			CHECK_CUDA(cudaMemcpy(d_z, z, nS*sizeof(float), cudaMemcpyHostToDevice));
			CHECK_CUDA_STATUS(cublasSgemv(cublashandle, CUBLAS_OP_T, nS, nDim, &alphaf, d_U, nS, d_z, 1, &betaf, d_dif, 1));
			CHECK_CUDA(cudaMemcpy(dif, d_dif, nDim*sizeof(float), cudaMemcpyDeviceToHost));

			tmp = 0.0f;
			for (int j = 0; j < nDim; j++)
				tmp += dif[j] * dif[j];
			gz = tmp / 2;

			//s_tmp = z - v;
			for (int j = 0; j < nS; j++)
				s_tmp[j] = z[j] - v[j];
			//gvz = gv + dgv'*s_tmp+b*s_tmp'*s_tmp / 2;
			tmp = 0.0f;
			for (int j = 0; j < nS; j++)
				tmp += dgv[j] * s_tmp[j] + s_tmp[j] * s_tmp[j] * b / 2;
			gvz = gv + tmp;
			if (gz <= gvz){
				beta[i + 1] = b;
				memcpy(z0, z1, nS*sizeof(float));
				memcpy(z1, z, nS*sizeof(float));
				break;
			}
		}
		if (beta[i + 1] == 0){
			beta[i + 1] = b;
			memcpy(z0, z1, nS*sizeof(float));
			memcpy(z1, z, nS*sizeof(float));
		}
		delta[i + 2] = 0.5*(1 + sqrtf(1 + 4 * delta[i + 1] * delta[i + 1]));
		sum_abs_z1_z0 = 0.0f;
		for (int j = 0; j < nS; j++){
			sum_abs_z1_z0 += (z1[j] - z0[j]) > 0 ? (z1[j] - z0[j]) : -1 * (z1[j] - z0[j]);
		}
		if (sum_abs_z1_z0 <= 1e-4)
			break;
		
	}
	memcpy(val_one_line, z1, nS*sizeof(float));

	CHECK_CUDA(cudaFree(d_v));
	CHECK_CUDA(cudaFree(d_U));
	CHECK_CUDA(cudaFree(d_x));
	CHECK_CUDA(cudaFree(d_dif));
	CHECK_CUDA(cudaFree(d_dgv));
	CHECK_CUDA(cudaFree(d_z));
	CHECK_CUDA(cudaFree(d_sXs_tmp));
	CHECK_CUDA(cudaFree(d_sX1_tmp));

	free(dif);
	free(v);
	free(dgv);
	free(s_tmp);
	free(z);
	free(z1);
	free(z0);
	free(beta);
	free(delta);
}